from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.window import Window
from pyspark.ml.evaluation import RegressionEvaluator
from pyspark.ml.recommendation import ALS
import pandas as pd

# 创建 Spark 会话
spark = SparkSession.builder.appName("ALSExample").getOrCreate()

# 读取数据
movies = spark.read.csv("file:///home/annie/movie_2/data/movies.csv",header=True,inferSchema=True)
ratings = spark.read.csv("file:///home/annie/movie_2/data/ratings.csv",header=True,inferSchema=True)
tags = spark.read.csv("file:///home/annie/movie_2/data/tags.csv",header=True,inferSchema=True)
#连接数据并删除同名数据
movies_renamed = movies.withColumnRenamed("movieId", "movieId_movies")
data = ratings.join(movies_renamed, ratings.movieId == movies_renamed.movieId_movies, how='inner')

# 删除重命名后的movieId_movies列
data = data.drop("movieId_movies")

# 删除任何包含空值的行
cleaned_data = data.na.drop(subset=['userId', 'movieId', 'rating'])

#划分数据集和训练集
(training,test) = cleaned_data.randomSplit([0.8,0.2])

# 创建并训练 ALS 模型
als = ALS(maxIter=5, regParam=0.01, userCol="userId", itemCol="movieId", ratingCol="rating", coldStartStrategy="drop")
model = als.fit(training)

# 在测试数据上进行预测
predictions = model.transform(test)

# 计算均方根误差
evaluator = RegressionEvaluator(metricName="rmse", labelCol="rating", predictionCol="prediction")
rmse = evaluator.evaluate(predictions)
print("Root-mean-square error = " + str(rmse))

# 对特定用户进行推荐
user_id=1
userRecs = model.recommendForAllUsers(10)
#userRecs.show()
# 使用explode函数将recommendations列中的数组拆分成多行
exploded_userRecs = userRecs.select(col("userId"), explode(col("recommendations")).alias("recommendation"))

# 将电影标题加入到推荐结果中
userRecs_with_titles = exploded_userRecs.join(movies, exploded_userRecs.recommendation.movieId == movies.movieId, "inner")
# 筛选出特定用户的电影推荐
user_1 = userRecs_with_titles[userRecs_with_titles['userId'] == user_id]
# 选择需要展示的列
selected_user_1 = user_1.select("userId", "recommendation.movieId", "recommendation.rating", "title")

selected_user_1.show()

spark.stop()